# -*- coding: utf-8 -*-
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Optional, List, Tuple, Set
import itertools
import logging
import numpy as np

from cqlib.circuits.circuit import Circuit

logger = logging.getLogger(__name__)

# -----------------------------------------------------------------------------
# 校验工具（与你原版一致）
# -----------------------------------------------------------------------------
def validate_in_set(name: str, value: object, values: Set[object]) -> None:
    if value not in values:
        raise ValueError("{} must be one of '{}', was '{}'.".format(name, values, value))

def validate_min(name: str, value: float, minimum: float) -> None:
    if value < minimum:
        raise ValueError('{} must have value >= {}, was {}'.format(name, minimum, value))

def validate_min_exclusive(name: str, value: float, minimum: float) -> None:
    if value <= minimum:
        raise ValueError('{} must have value > {}, was {}'.format(name, minimum, value))

def validate_max(name: str, value: float, maximum: float) -> None:
    if value > maximum:
        raise ValueError('{} must have value <= {}, was {}'.format(name, maximum, value))

def validate_max_exclusive(name: str, value: float, maximum: float) -> None:
    if value >= maximum:
        raise ValueError('{} must have value < {}, was {}'.format(name, maximum, value))

def validate_range(name: str, value: float, minimum: float, maximum: float) -> None:
    if value < minimum or value > maximum:
        raise ValueError(
            '{} must have value >= {} and <= {}, was {}'.format(name, minimum, maximum, value))

def validate_range_exclusive(name: str, value: float, minimum: float, maximum: float) -> None:
    if value <= minimum or value >= maximum:
        raise ValueError(
            '{} must have value > {} and < {}, was {}'.format(name, minimum, maximum, value))

def validate_range_exclusive_min(name: str, value: float,
                                 minimum: float, maximum: float) -> None:
    if value <= minimum or value > maximum:
        raise ValueError(
            '{} must have value > {} and <= {}, was {}'.format(name, minimum, maximum, value))

def validate_range_exclusive_max(name: str, value: float,
                                 minimum: float, maximum: float) -> None:
    if value < minimum or value >= maximum:
        raise ValueError(
            '{} must have value >= {} and < {}, was {}'.format(name, minimum, maximum, value))


# # -----------------------------------------------------------------------------
# # 抽象基类（接口与原版一致；实现改为 cqlib 语义）
# # -----------------------------------------------------------------------------
# class Reciprocal(ABC):
#     @abstractmethod
#     def __init__(self) -> None:
#         pass

#     @abstractmethod
#     def construct_circuit(self, mode: str, inreg: List[int], circuit: Optional[Circuit] = None) -> Circuit:
#         raise NotImplementedError()

# -----------------------------------------------------------------------------
# cqlib gate helpers
# -----------------------------------------------------------------------------
def _has(circ: Circuit, name: str) -> bool:
    return hasattr(circ, name)

def _mcx_chain(circ: Circuit, controls: List[int], target: int, aux: List[int]):
    """多控 X：Toffoli 链式分解，需要 len(controls)-2 个 ancilla。"""
    k = len(controls)
    if k == 0:
        circ.x(target); return
    if k == 1:
        circ.cx(controls[0], target); return
    if k == 2:
        circ.ccx(controls[0], controls[1], target); return
    need = k - 2
    if len(aux) < need:
        raise ValueError(f"MCX needs {need} ancillas, only {len(aux)} provided.")
    circ.ccx(controls[0], controls[1], aux[0])
    for i in range(2, k - 1):
        circ.ccx(controls[i], aux[i - 2], aux[i - 1])
    circ.ccx(controls[k - 1], aux[need - 1], target)
    for i in reversed(range(2, k - 1)):
        circ.ccx(controls[i], aux[i - 2], aux[i - 1])
    circ.ccx(controls[0], controls[1], aux[0])

def _mcry(circ: Circuit, theta: float, controls: List[int], target: int, aux: List[int]):
    """
    多控 RY：优先用 CRY（AND 所有控制到单 ancilla），
    ancilla 不够时回退到 MCX+RY/−RY 分解。
    """
    k = len(controls)
    if k == 0:
        circ.ry(target,theta); return
    if k == 1:
        circ.cry(controls[0], target, theta); return

    need_ctrl_and = k - 1   # CRY 路径所需 ancilla
    need_chain    = k - 2   # 回退路径所需 ancilla

    if _has(circ, "cry") and len(aux) >= need_ctrl_and:
        # ---- CRY 路径：把所有控制位与到单 ancilla ----
        a = aux[:need_ctrl_and]
        circ.ccx(controls[0], controls[1], a[0])
        for i in range(2, k):
            circ.ccx(controls[i], a[i - 2], a[i - 1])
        # a[k-2] = AND(controls)
        circ.cry(a[k - 2], target, theta)
        # 反向解算
        for i in reversed(range(2, k)):
            circ.ccx(controls[i], a[i - 2], a[i - 1])
        circ.ccx(controls[0], controls[1], a[0])
    else:
        # ---- 回退：MCX + RY/−RY ----
        _mcx_chain(circ, controls, target, aux)
        circ.ry(target, theta / 2.0)
        _mcx_chain(circ, controls, target, aux)
        circ.ry(target, -theta / 2.0)


# -----------------------------------------------------------------------------
# LookupRotation（cqlib 版本，已切到 _mcry）
# -----------------------------------------------------------------------------
class LookupRotation(ABC):
    def __init__(
        self,
        pat_length: Optional[int] = None,
        subpat_length: Optional[int] = None,
        scale: float = 0.0,
        negative_evals: bool = False,
        evo_time: Optional[float] = None,
        lambda_min: Optional[float] = None,
    ) -> None:
        validate_range("scale", scale, 0, 1)
        super().__init__()
        self._pat_length = pat_length
        self._subpat_length = subpat_length
        self._negative_evals = negative_evals
        self._scale = scale
        self._evo_time = evo_time
        self._lambda_min = lambda_min

        # 运行期上下文（索引）
        self._ev_idx: List[int] = []
        self._work: Optional[int] = None
        self._msq: Optional[int] = None
        self._anc: Optional[int] = None
        self._aux_pool: List[int] = []

        self._reg_size = 0
        self._circuit: Optional[Circuit] = None

    @staticmethod
    def _classic_approx(k, n, m, negative_evals=False):
        def bin_to_num(binary):
            return np.sum([2 ** -(n + 1) for n, i in enumerate(reversed(binary)) if i == '1'])
        def get_est_lamb(vec, fo, n, k):
            if fo - n > 0:
                remainder = sum([2 ** -i for i in range(k - (fo - n - 1), k + 1)])
                return bin_to_num(vec) + remainder / 2
            return bin_to_num(vec)

        from collections import OrderedDict
        output = OrderedDict()
        fo = None
        for fo in range(k - 1, n - 1, -1):
            if negative_evals and fo == k - 1:
                continue
            vec = ['0'] * k
            vec[fo] = '1'
            for pattern_ in itertools.product('10', repeat=m):
                app_pattern_array = []
                lambda_array = []
                for appendpat in itertools.product('10', repeat=n - m):
                    pattern = pattern_ + appendpat
                    vec[fo - n:fo] = pattern
                    e_l = get_est_lamb(vec.copy(), fo, n, k)
                    lambda_array.append(e_l)
                    app_pattern_array.append(list(reversed(appendpat)))
                fo_pos = k - fo - 1
                prev_res = output.get(fo_pos, [])
                output[fo_pos] = prev_res + [(list(reversed(pattern_)),
                                              app_pattern_array, lambda_array)]
        last_fo = fo
        vec = ['0'] * k
        for pattern_ in itertools.product('10', repeat=m):
            app_pattern_array = []
            lambda_array = []
            for appendpat in itertools.product('10', repeat=n - m):
                pattern = list(pattern_ + appendpat).copy()
                if '1' not in pattern and (not negative_evals):
                    continue
                if '1' not in pattern and negative_evals:
                    e_l = 0.5
                else:
                    vec[last_fo - n:last_fo] = list(pattern)
                    e_l = get_est_lamb(vec.copy(), last_fo, n, k)
                lambda_array.append(e_l)
                app_pattern_array.append(list(reversed(appendpat)))
            fo_pos = k - last_fo
            prev_res = output.get(fo_pos, [])
            output[fo_pos] = prev_res + [(list(reversed(pattern_)),
                                          app_pattern_array, lambda_array)]
        return output

    def _set_msq(self, qc: Circuit, ev_reg: List[int], fo_pos: int, last_iteration: bool = False):
        ev = list(ev_reg)
        msq = self._msq; assert msq is not None
        if last_iteration:
            if fo_pos == 1:
                qc.x(ev[0]); qc.cx(ev[0], msq); qc.x(ev[0])
            elif fo_pos == 2:
                qc.x(ev[0]); qc.x(ev[1]); qc.ccx(ev[0], ev[1], msq); qc.x(ev[1]); qc.x(ev[0])
            elif fo_pos > 2:
                for idx in range(fo_pos): qc.x(ev[idx])
                _mcx_chain(qc, ev[:fo_pos], msq, self._aux_pool)
                for idx in range(fo_pos): qc.x(ev[idx])
            else:
                qc.x(msq)
        elif fo_pos == 0:
            qc.cx(ev[0], msq)
        elif fo_pos == 1:
            qc.x(ev[0]); qc.ccx(ev[0], ev[1], msq); qc.x(ev[0])
        elif fo_pos > 1:
            for idx in range(fo_pos): qc.x(ev[idx])
            _mcx_chain(qc, ev[:fo_pos + 1], msq, self._aux_pool)
            for idx in range(fo_pos): qc.x(ev[idx])
        else:
            raise RuntimeError("first-one register index < 0")

    def _set_bit_pattern(self, qc: Circuit, pattern: List[str], tgt: int, offset: int):
        for c, bit in enumerate(pattern):
            if bit == '0':
                qc.x(self._ev_idx[int(c + offset)])
        L = len(pattern)
        if L > 2:
            ctrls = [self._ev_idx[i] for i in range(offset, offset + L)]
            _mcx_chain(qc, ctrls, tgt, self._aux_pool)
        elif L == 2:
            qc.ccx(self._ev_idx[offset], self._ev_idx[offset + 1], tgt)
        elif L == 1:
            qc.cx(self._ev_idx[offset], tgt)
        for c, bit in enumerate(pattern):
            if bit == '0':
                qc.x(self._ev_idx[int(c + offset)])

    def construct_circuit(self, mode: str, inreg: List[int], circuit: Optional[Circuit] = None) -> Circuit:
        """
        Args:
            mode: 仅支持 'circuit'
            inreg: 特征值寄存器在目标电路中的索引列表（长度 k）
            circuit: None=新建子电路；否则在现有电路后追加
        """
        if mode != "circuit":
            raise NotImplementedError("Only 'circuit' mode is supported in cqlib version.")

        k = len(inreg)
        self._reg_size = k

        if self._lambda_min and self._evo_time:
            self._scale = self._lambda_min / (2 * np.pi) * self._evo_time
        if self._scale == 0:
            self._scale = 2 ** -k

        if self._pat_length is None:
            self._pat_length = (k - (2 if self._negative_evals else 1)) if k <= 6 else 5
        if k <= self._pat_length:
            self._pat_length = k - (2 if self._negative_evals else 1)
        if self._subpat_length is None:
            self._subpat_length = int(np.ceil(self._pat_length / 2))
        n = self._pat_length
        m = self._subpat_length

        # 估算最坏情况下的 ancilla 需求：CRY 路径需要 k-1
        aux_needed = max(0, (k - 1))

        if circuit is None:
            total = k + 3 + aux_needed
            qc = Circuit(total)
            self._ev_idx = list(range(0, k))
            self._work = k
            self._msq = k + 1
            self._anc = k + 2
            self._aux_pool = list(range(k + 3, k + 3 + aux_needed))
        else:
            qc = circuit
            self._ev_idx = list(inreg)
            start = max(inreg) + 1
            self._work = start
            self._msq = start + 1
            self._anc = start + 2
            self._aux_pool = list(range(start + 3, start + 3 + aux_needed))

        assert self._work is not None and self._msq is not None and self._anc is not None

        approx_dict = self._classic_approx(k, n, m, negative_evals=self._negative_evals)

        fo = None
        old_fo = None
        ev_all = list(self._ev_idx)
        ev_wo_sign = ev_all[1:] if self._negative_evals else ev_all

        for fo in list(approx_dict.keys()):
            pattern_map = approx_dict[fo]

            if self._negative_evals:
                if old_fo != fo:
                    if old_fo is not None:
                        self._set_msq(qc, ev_wo_sign, int(old_fo - 1))
                    old_fo = fo
                    if fo + n == k:
                        self._set_msq(qc, ev_wo_sign, int(fo - 1), last_iteration=True)
                    else:
                        self._set_msq(qc, ev_wo_sign, int(fo - 1), last_iteration=False)
            else:
                if old_fo != fo:
                    if old_fo is not None:
                        self._set_msq(qc, ev_all, int(old_fo))
                    old_fo = fo
                    if fo + n == k:
                        self._set_msq(qc, ev_all, int(fo), last_iteration=True)
                    else:
                        self._set_msq(qc, ev_all, int(fo), last_iteration=False)

            offset_mpat = fo + (n - m) if fo < k - n else fo + n - m - 1
            for mainpat, subpat_list, lambda_arr in pattern_map:
                self._set_bit_pattern(qc, mainpat, self._work, offset_mpat + 1)

                for subpattern, lambd in zip(subpat_list, lambda_arr):
                    theta = 2 * np.arcsin(min(1.0, self._scale / lambd))
                    offset = fo + 1 if fo < k - n else fo

                    # —— 关键处：两次多控旋转，已切换为 _mcry（优先 CRY 路径） ——
                    _mcry(qc, theta / 2.0, [self._work, self._msq], self._anc, self._aux_pool)
                    self._set_bit_pattern(qc, subpattern, self._anc, offset)
                    _mcry(qc, -theta / 2.0, [self._work, self._msq], self._anc, self._aux_pool)
                    self._set_bit_pattern(qc, subpattern, self._anc, offset)

                self._set_bit_pattern(qc, mainpat, self._work, offset_mpat + 1)

        last_fo = fo
        if self._negative_evals:
            self._set_msq(qc, ev_wo_sign, int(last_fo - 1), last_iteration=True)
        else:
            self._set_msq(qc, ev_all, int(last_fo), last_iteration=True)

        if self._negative_evals:
            # qc.cry(self._ev_idx[0], self._anc, 2 * np.pi)
            # qc.cry(self._anc, self._work, -2 * np.pi)
            qc.cry(self._ev_idx[0], self._anc, 2*np.pi)

        self._circuit = qc
        return qc


